import os

import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms


class BrainMRIDataset(Dataset):
    def __init__(self, base_dir, transform=None):
        """
        初始化数据集
        :param base_dir: 包含所有病人文件夹的根目录
        :param transform: 应用于图像的转换
        """
        self.base_dir = base_dir
        self.transform = transform
        self.samples = []

        # 遍历每个病人的文件夹
        patient_dirs = [
            os.path.join(base_dir, name)
            for name in os.listdir(base_dir)
            if os.path.isdir(os.path.join(base_dir, name))
        ]
        for patient_dir in patient_dirs:
            images = sorted(
                [
                    file
                    for file in os.listdir(patient_dir)
                    if file.endswith(".tif") and not file.endswith("_mask.tif")
                ]
            )
            masks = {
                os.path.splitext(os.path.basename(mask))[0]: os.path.join(
                    patient_dir, mask
                )
                for mask in sorted(
                    [
                        file
                        for file in os.listdir(patient_dir)
                        if file.endswith("_mask.tif")
                    ]
                )
            }

            images_paths = [
                os.path.join(patient_dir, image) for image in images
            ]  # 正确定义images_paths

            for image_path in images_paths:
                image_base = os.path.splitext(os.path.basename(image_path))[0]
                mask_base = f"{image_base}_mask"
                if mask_base in masks:
                    self.samples.append((image_path, masks[mask_base]))
                else:
                    raise ValueError(f"No matching mask found for image: {image_path}")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        image_path, mask_path = self.samples[idx]
        image = Image.open(image_path).convert("L")
        mask = Image.open(mask_path).convert("L")

        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)

        return image, mask


# 定义转换
transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])]
)

# 加载数据集
data_dir = "lgg-mri-segmentation/kaggle_3m"  # 请确保此路径正确指向您的数据文件夹
full_dataset = BrainMRIDataset(base_dir=data_dir, transform=transform)

# 分割数据集
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(
    full_dataset, [train_size, val_size]
)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)

print(f"Total dataset size: {len(full_dataset)}")
print(f"Train dataset size: {len(train_dataset)}")
print(f"Validation dataset size: {len(val_dataset)}")
